#include "coro/coro.hpp"
#include "pw_allocator/first_fit.h"
#include <zephyr/fff.h>
#include <zephyr/ztest.h>

DEFINE_FFF_GLOBALS;

FAKE_VOID_FUNC(func1);
FAKE_VOID_FUNC(func2);
FAKE_VOID_FUNC(func3);

namespace {
struct coro_fixture {
  static auto Setup() -> void * {
    static coro_fixture f;
    // Set up code for test suite
    return &f;
  }
  static auto Before(void *f) -> void {
    func1_fake.call_count = 0;
    func2_fake.call_count = 0;
    func3_fake.call_count = 0;
  };
  static auto After(void *f) -> void {};
  static auto Teardown(void *f) -> void {};
  CO_EVT_DECLARE(e1, int);
  CO_EVT_DECLARE(e2, int, int);
  CO_EVT_DECLARE(e3);
  CORO_MEM_CTX(m_ctx, 1024);
};
using Fixture = coro_fixture;
} // namespace

ZTEST_F(coro, test_event) {
  auto t1 = coro::Coro([&](CO_CTX) -> coro::AsyncStatus {
    co_await *fixture->e3;
    func1();
    co_return coro::OkStatus();
  }(fixture->m_ctx));
  auto t2 = coro::Coro([&](CO_CTX) -> coro::AsyncStatus {
    co_await *fixture->e3;
    func2();
    co_return coro::OkStatus();
  }(fixture->m_ctx));
  auto t3 = coro::Coro([&](CO_CTX) -> coro::AsyncStatus {
    co_await *fixture->e3;
    func3();
    co_return coro::OkStatus();
  }(fixture->m_ctx));
  coro::Dispatcher d;
  d.Post(t1, t2, t3);
  zassert_true(d.RunUntilStalled().IsPending());
  fixture->e3();
  zassert_false(d.RunUntilStalled().IsPending());
  zassert_equal(func1_fake.call_count, 1);
  zassert_equal(func2_fake.call_count, 1);
  zassert_equal(func3_fake.call_count, 1);
}

ZTEST_F(coro, test_when_any) {
  auto t1 = [&](coro::Context &) -> coro::AsyncStatus {
    auto [r1, r2] =
        co_await coro::WhenAny(fixture->e1.Get(), fixture->e2.Get());
    zassert_true(r1.has_value());
    zassert_equal(666, r1.value());
    zassert_false(r2.has_value());
    co_return pw::OkStatus();
  };
  auto t2 = [&](coro::Context &) -> coro::AsyncStatus {
    auto [r1, r2] =
        co_await coro::WhenAny(fixture->e1.Get(), fixture->e2.Get());
    zassert_false(r1.has_value());
    zassert_true(r2.has_value());
    const auto &[v1, v2] = r2.value();
    zassert_equal(v1, 111);
    zassert_equal(v2, 222);
    co_return pw::OkStatus();
  };
  auto c1 = coro::Coro(t1(fixture->m_ctx));
  auto c2 = coro::Coro(t2(fixture->m_ctx));
  coro::Dispatcher d;
  zassert_true(d.RunPendableUntilStalled(c1).IsPending());
  fixture->e1.Set(666);
  zassert_true(d.RunPendableUntilStalled(c1).IsReady());

  zassert_true(d.RunPendableUntilStalled(c2).IsPending());
  fixture->e2.Set(111, 222);
  zassert_true(d.RunPendableUntilStalled(c2).IsReady());
}

ZTEST_SUITE(coro, NULL, Fixture::Setup, Fixture::Before, Fixture::After,
            Fixture::Teardown);